import pygame
import random

# Window size
WIDTH, HEIGHT = 800, 600

# Colors
WHITE = (255, 255, 255)
BLACK = (0, 0, 0)

class Boid:
    def __init__(self):
        self.x = random.randint(0, WIDTH)
        self.y = random.randint(0, HEIGHT)
        self.velocity_x = random.uniform(-1.5, 1.5)
        self.velocity_y = random.uniform(-1.5, 1.5)

    def cohesion(self, boids):
        total_x = 0
        total_y = 0

        for boid in boids:
            if boid != self:
                dx = boid.x - self.x
                dy = boid.y - self.y

                distance = (dx**2 + dy**2)**0.5

                if distance > 0:
                    total_x += dx / distance
                    total_y += dy / distance

        avg_x = total_x / len(boids)
        avg_y = total_y / len(boids)

        self.velocity_x += avg_x * 0.01
        self.velocity_y += avg_y * 0.01

    def separation(self, boids):
        repulsion = 0
        for boid in boids:
            if boid != self:
                dx = boid.x - self.x
                dy = boid.y - self.y

                distance = (dx**2 + dy**2)**0.5

                if distance < 50:
                    repulsion += (1 / distance) * (-dx, -dy)

        self.velocity_x += repulsion[0] * 0.01
        self.velocity_y += repulsion[1] * 0.01

    def alignment(self, boids):
        avg_velocity = (0, 0)
        for boid in boids:
            if boid != self:
                avg_velocity = tuple((avg_velocity[0] + boid.velocity_x,
                                       avg_velocity[1] + boid.velocity_y))

        avg_velocity = (avg_velocity[0] / len(boids),
                        avg_velocity[1] / len(boids))

        diff_x = avg_velocity[0] - self.velocity_x
        diff_y = avg_velocity[1] - self.velocity_y

        if abs(diff_x) < 0.5 and abs(diff_y) < 0.5:
            self.velocity_x += diff_x * 0.01
            self.velocity_y += diff_y * 0.01

    def update(self):
        self.x += self.velocity_x
        self.y += self.velocity_y

        if self.x - 50 < 0 or self.x + 50 > WIDTH:
            self.velocity_x *= -1
        if self.y - 50 < 0 or self.y + 50 > HEIGHT:
            self.velocity_y *= -1

def main():
    pygame.init()
    screen = pygame.display.set_mode((WIDTH, HEIGHT))
    clock = pygame.time.Clock()

    boids = [Boid() for _ in range(100)]

    running = True
    while running:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                running = False

        screen.fill(BLACK)

        for i, boid in enumerate(boids):
            boid.cohesion(boids)
            boid.separation(boids)
            boid.alignment(boids)
            boid.update()

            pygame.draw.circle(screen, WHITE, (int(boid.x), int(boid.y)), 5)

        pygame.display.flip()
        clock.tick(60)

    pygame.quit()

if __name__ == "__main__":
    main()